# 测试脚本

import json
import joblib
from typing import Dict, Any, List, Tuple, Optional
import pandas as pd
from prepare_ali_data import (
    calculate_experience_score,
    calculate_education_score,
    calculate_compensation_score,
    calculate_responsibility_score,
    determine_location_tier,
    determine_tech_level,
    determine_role_type,
    determine_school_level
)

VALID_SCHOOLS = {
    '985': [
        '清华大学', '北京大学', '浙江大学', '上海交通大学', '复旦大学',
        '南京大学', '中国科学技术大学', '哈尔滨工业大学', '西安交通大学',
        '华中科技大学', '武汉大学', '吉林大学', '四川大学', '山东大学',
        '同济大学', '华南理工大学', '中山大学', '北京理工大学',
        '天津大学', '东南大学', '南开大学', '厦门大学', '大连理工大学',
        '华东师范大学', '中南大学', '湖南大学', '北京航空航天大学',
        '重庆大学', '兰州大学', '电子科技大学', '东北大学',
        '中国人民大学', '北京师范大学', '中国农业大学', '中央民族大学',
        '西北工业大学', '西北农林科技大学', '中国海洋大学'
    ],
    '211': [
        '清华大学', '北京大学', '浙江大学', '上海交通大学', '复旦大学',
        '南京大学', '中国科学技术大学', '哈尔滨工业大学', '西安交通大学',
        '华中科技大学', '武汉大学', '吉林大学', '四川大学', '山东大学',
        '同济大学', '华南理工大学', '中山大学', '北京理工大学',
        '天津大学', '东南大学', '南开大学', '厦门大学', '大连理工大学',
        '华东师范大学', '中南大学', '湖南大学', '北京航空航天大学',
        '重庆大学', '兰州大学', '电子科技大学', '东北大学',
        '中国人民大学', '北京师范大学', '中国农业大学', '中央民族大学',
        '西北工业大学', '西北农林科技大学', '中国海洋大学',
        '北京邮电大学', '北京工业大学', '北京科技大学', '北京化工大学',
        '北京交通大学', '华北电力大学', '北京林业大学', '北京中医药大学',
        '北京协和医学院', '北京外国语大学', '中国政法大学', '首都医科大学',
        '南京理工大学', '南京航空航天大学', '苏州大学', '南京农业大学',
        '中国药科大学', '华东理工大学', '东华大学', '上海大学',
        '上海财经大学', '上海外国语大学', '暨南大学', '华南师范大学',
        '郑州大学', '西北大学', '西南交通大学', '广西大学', '云南大学',
        '贵州大学', '新疆大学', '石河子大学', '西藏大学', '宁夏大学',
        '青海大学', '内蒙古大学', '海南大学', '河北工业大学', '太原理工大学',
        '延边大学', '辽宁大学', '大连海事大学', '东北师范大学', '哈尔滨工程大学',
        '东北林业大学', '东北农业大学', '福州大学', '合肥工业大学', '安徽大学',
        '南昌大学', '河南大学', '武汉理工大学', '华中师范大学', '华中农业大学',
        '中南财经政法大学', '湖南师范大学', '国防科技大学', '四川农业大学',
        '西南财经大学', '重庆大学', '西南大学', '贵州大学', '云南大学',
        '西北大学', '西安电子科技大学', '陕西师范大学', '空军军医大学', '长安大学',
        '中国矿业大学', '中国地质大学', '河海大学', '江南大学', '南京师范大学',
        '中国石油大学', '中国海洋大学', '北京交通大学', '北京工业大学', '北京科技大学'
    ]
}

COMMON_TITLES = {
    '开发': [
        '初级开发工程师', '开发工程师', '高级开发工程师',
        '资深开发工程师', '技术专家', '高级技术专家',
        '研发工程师', '高级研发工程师', '资深研发工程师'
    ],
    '产品经理': [
        '产品经理', '高级产品经理', '资深产品经理',
        '产品专家', '高级产品专家'
    ],
    '主管': [
        '技术经理', '高级技术经理', '技术总监',
        '研发经理', '研发总监', '技术负责人'
    ],
    '设计': [
        '设计师', '高级设计师', '资深设计师',
        '设计专家', '高级设计专家'
    ]
}

def normalize_school(school: str) -> str:
    """标准化学校名称"""
    if not school or school == '无':
        return '其他'

    school = school.replace('大学', '').replace('学院', '')

    for level, schools in VALID_SCHOOLS.items():
        for valid_school in schools:
            if school in valid_school or valid_school in school:
                return level
                
    return '其他'

def normalize_title(title: str) -> str:
    """标准化职位名称"""
    if not title:
        return '其他'
        
    title = title.lower()
    
    for category, titles in COMMON_TITLES.items():
        for valid_title in titles:
            if title == valid_title.lower():
                return valid_title
                
    if 'java' in title:
        return normalize_title('开发工程师')
    if '算法' in title:
        return normalize_title('算法工程师')
    if '架构' in title:
        return normalize_title('架构师')
        
    return '开发工程师'  # 默认职位

def get_valid_input(prompt: str, valid_options: Optional[List[str]] = None, is_number: bool = False) -> str:
    """获取有效的用户输入"""
    while True:
        value = input(prompt).strip()
        
        if not value:
            print("输入不能为空，请重新输入！")
            continue
            
        if is_number:
            try:
                float(value)
                return value
            except ValueError:
                print("请输入有效的数字！")
                continue
                
        if valid_options:
            if value not in valid_options:
                print(f"请输入以下选项之一: {', '.join(valid_options)}")
                continue
                
        return value

def get_user_features() -> Dict[str, Any]:
    """获取用户输入的特征"""
    print("\n" + "="*50)
    print("职级预测简化版")
    print("="*50 + "\n")
    
    print("\n--- 工作经验 ---")
    features = {}
    
    features['year_of_experience'] = float(get_valid_input(
        "请输入您的总工作年限（年）: ",
        is_number=True
    ))
    
    features['year_in_company'] = float(get_valid_input(
        "请输入您在当前公司的工作年限（年）: ",
        is_number=True
    ))
    
    print("\n--- 教育背景 ---")
    print("提示：支持的学校包括所有985/211高校")
    features['degree'] = get_valid_input(
        "请选择您的最高学历 [专科, 本科, 硕士, 博士]: ",
        valid_options=['专科', '本科', '硕士', '博士']
    )
    
    features['bachelor_school'] = normalize_school(get_valid_input(
        "请输入您的本科学校（如果没有本科学历，请输入'无'）: "
    ))
    
    if features['degree'] in ['硕士', '博士']:
        features['advanced_school'] = normalize_school(get_valid_input(
            "请输入您的研究生学校: "
        ))
    else:
        features['advanced_school'] = '其他'
    
    print("\n--- 工作信息 ---")

    print("职位类型选择：")
    job_types = ['开发', '产品经理', '主管', '设计']
    features['job_type'] = get_valid_input(
        "请选择您的职位类型 [开发, 产品经理, 主管, 设计]: ",
        valid_options=job_types
    )
    
    # 根据职位类型设置对应的标准化职位名称
    job_type_mapping = {
        '开发': '开发工程师',
        '产品经理': '产品经理', 
        '主管': '技术经理',
        '设计': '设计师'
    }
    features['title'] = job_type_mapping[features['job_type']]
    
    print("\n城市等级说明：")
    print("一线城市: 北京、上海、广州、深圳")
    print("新一线城市: 杭州、南京、武汉、成都、西安")
    print("其他城市: 除上述外的所有城市")
    
    city_input_type = get_valid_input(
        "\n请选择输入方式 [1-从列表选择, 2-直接输入城市名]: ",
        valid_options=['1', '2']
    )
    
    if city_input_type == '1':
        tier1_cities = ['北京', '上海', '广州', '深圳']
        # 成都、杭州、重庆、武汉、苏州、西安、 南京、长沙、郑州、天津、合肥、青岛、东莞、宁波与佛山
        new_tier1_cities = ['杭州', '南京', '武汉', '成都', '西安', '重庆', '苏州', '长沙', '郑州', '天津', '合肥', '青岛', '东莞', '宁波', '佛山']
        all_cities = tier1_cities + new_tier1_cities + ['其他城市']
        
        print("\n城市列表：")
        for i, city in enumerate(all_cities, 1):
            print(f"{i}. {city}")
            
        while True:
            try:
                choice = int(get_valid_input("请输入城市编号: ", is_number=True))
                if 1 <= choice <= len(all_cities):
                    if choice == len(all_cities):  # 选择了"其他城市"
                        features['location'] = get_valid_input("请输入您的城市名称: ")
                    else:
                        features['location'] = all_cities[choice-1]
                    break
                else:
                    print(f"请输入 1-{len(all_cities)} 之间的数字！")
            except ValueError:
                print("请输入有效的数字！")
    else:
        features['location'] = get_valid_input(
            "请输入您工作的城市: "
        )
        
        from prepare_ali_data import determine_location_tier
        tier = determine_location_tier(features['location'])
        print(f"您输入的城市 '{features['location']}' 被识别为: {tier}")
    
    print("\n--- 薪资信息（单位：元/年）---")
    features['total_comp'] = float(get_valid_input(
        "请输入您的年薪（包含五险一金）: ",
        is_number=True
    ))
    
    # 计算月薪
    features['base_monthly_comp'] = features['total_comp'] / 12
    
    features['bonus_comp'] = 0  # 年薪已包含奖金，不再单独计算
    
    features['stock_comp'] = 0
    
    return features

def prepare_features(user_input: Dict[str, Any]) -> Dict[str, Any]:
    """准备模型所需的特征"""
    features = {}
    
    # 计算各维度评分
    features['experience_score'] = calculate_experience_score(
        user_input['year_of_experience'],
        user_input['year_in_company']
    )
    
    features['education_score'] = calculate_education_score(
        user_input['degree'],
        determine_school_level(
            user_input['bachelor_school'],
            user_input['advanced_school']
        )
    )
    
    features['technical_score'] = 0 
    features['workload_score'] = 70  
    
    features['compensation_score'] = calculate_compensation_score(
        user_input['total_comp'],  
        user_input['location'],
        user_input['base_monthly_comp'],  
        user_input['bonus_comp'],  
        0  
    )
    
    features['responsibility_score'] = calculate_responsibility_score(
        user_input['title'],
        '技术部',
        user_input['total_comp']
    )
    
    # 分类特征
    features['location_tier'] = determine_location_tier(
        user_input['location']
    )
    
    features['tech_level'] = determine_tech_level(
        user_input['title'],
        user_input['year_of_experience']
    )
    
    # 根据职位类型直接决定角色类型
    job_type_to_role = {
        '开发工程师': 'IC',
        '产品经理': 'IC', 
        '技术经理': '管理',
        '设计师': 'IC'
    }
    features['role_type'] = job_type_to_role.get(user_input['title'], 'IC')
    
    features['business_domain'] = '其他'
    
    return features

def format_prediction_result(
    predicted_level: str,
    confidence: float,
    probabilities: Dict[str, float],
    feature_importance: Dict[str, float]
) -> str:
    """格式化预测结果"""
    result = [
        "\n" + "="*50,
        "预测结果",
        "="*50,
        f"\n预测职级: {predicted_level}",
        f"置信度: {confidence:.2%}",
        "\n各职级概率:",
    ]
    
    # 按职级排序显示概率
    for level, prob in sorted(probabilities.items()):
        result.append(f"{level}: {prob:.2%}")
    
    result.extend([
        "\n影响因素分析（按重要性排序）:",
        "-"*30
    ])
    
    score_features = {
        'compensation_score': '薪酬评分',
        'responsibility_score': '责任评分',
        'experience_score': '经验评分',
        'education_score': '教育评分'
    }
    
    # 按重要性排序输出评分特征
    sorted_features = sorted(
        [(k, v) for k, v in feature_importance.items() if k in score_features],
        key=lambda x: x[1],
        reverse=True
    )
    
    for feature, importance in sorted_features:
        result.append(
            f"{score_features[feature]}: {importance:.2%}"
        )
    
    return "\n".join(result)

def main():
    try:
        model = joblib.load('model_simple.joblib')
    except FileNotFoundError:
        print("错误：找不到模型文件")
        return
    
    user_input = get_user_features()
    
    features = prepare_features(user_input)
    
    predicted_level, confidence, probabilities = model.predict_level(features)
    
    # 获取特征重要性
    feature_importance = {}
    if hasattr(model.model, 'named_steps') and 'classifier' in model.model.named_steps:
        classifier = model.model.named_steps['classifier']
        if hasattr(classifier, 'feature_importances_'):
            for feature, score in zip(model.numeric_features + model.categorical_features, 
                                    classifier.feature_importances_):
                feature_importance[feature] = score
    
    if not feature_importance:
        feature_importance = {
            'compensation_score': 0.35,
            'responsibility_score': 0.25,
            'experience_score': 0.25,
            'education_score': 0.15
        }
    
    print(format_prediction_result(
        predicted_level,
        confidence,
        probabilities,
        feature_importance
    ))

if __name__ == '__main__':
    main()